package com.woniu.util;

import cn.hutool.json.JSONUtil;
import lombok.Data;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.stream.Collectors;

public class TreeUtil {

    /**
     * listToTree
     * @param target      需转换的数据
     * @param getId       主键
     * @param getParentId 父id （父id必须和主键相同类型）
     * @param getChildren 子集
     * @param setChildren 子集
     * @return tree
     * 使用方法：
     * 1.实体类配置除子集外的构造方法;
     * 2.接收 List<实体类> 实体类List = TreeUtil.listToTree(list(需转换的集合), TreeDto(实体类)::getId(主键的get方法), TreeDto(实体类)::getPid((父ID的get方法)),
     *                 (实体类)TreeDto::getChildren(下一级的get方法), (实体类)TreeDto::setChildren(下一级的set方法));
     * 3.转成Json数据返回 JSONUtil.toJsonStr(实体类List);
     * 示例见下
     */
    public static <T, R> List<T> listToTree(List<T> target, Function<T, R> getId, Function<T, R> getParentId,
                                            Function<T, List<T>> getChildren, BiConsumer<T, List<T>> setChildren) {

        Map<R, T> oldMap = target.stream().collect(Collectors.toMap(getId, T -> T));
        List<T> result = new ArrayList<>();
        target.forEach(tree -> {
            T parent = oldMap.get(getParentId.apply(tree));
            if (parent == null) {
                result.add(tree);
            } else {
                List<T> ch = getChildren.apply(parent);
                if (ch == null) {
                    ch = new ArrayList<>();
                }
                ch.add(tree);
                setChildren.accept(parent, ch);
            }
        });
        return result;
    }


    public static void main(String[] args) {
        //模拟数据
        List<TreeDto> list = new ArrayList<TreeDto>() {{
            add(new TreeDto("1", "0", "第1层"));
            add(new TreeDto("11", "1", "第2层"));
            add(new TreeDto("12", "1", "第2层"));
            add(new TreeDto("2", "0", "第1层"));
            add(new TreeDto("3", "4", "第4层"));
            add(new TreeDto("4", "11", "第3层"));
        }};

        System.out.println(JSONUtil.toJsonStr(TreeUtil.listToTree(list, TreeDto::getId, TreeDto::getPid,
                TreeDto::getChildren, TreeDto::setChildren)));
    }

    @Data
    static class TreeDto {
        // 除子集外的构造方法
        public TreeDto(String id, String pid, String name) {
            this.id = id;
            this.pid = pid;
            this.name = name;
        }

        private String id;
        private String pid;//父id
        private String name;
        private List<TreeDto> children;

    }

}

